package yycore

import chisel3._
import chisel3.util._
import common.Constants._

object ALUOpType {
  // ALU Operation Signal
  val X   = 0.asUInt(fuOpType.W)
  val ADD = 0.asUInt(fuOpType.W)
  val SUB = 1.asUInt(fuOpType.W)
  val SLL = 2.asUInt(fuOpType.W)
  val SRL = 3.asUInt(fuOpType.W)
  val SRA = 4.asUInt(fuOpType.W)
  val AND = 5.asUInt(fuOpType.W)
  val OR  = 6.asUInt(fuOpType.W)
  val XOR = 7.asUInt(fuOpType.W)
  val SLT = 8.asUInt(fuOpType.W)
  val SLTU= 9.asUInt(fuOpType.W)

  val ADDW = 10.asUInt(fuOpType.W)
  val SUBW = 11.asUInt(fuOpType.W)
  val SLLW = 12.asUInt(fuOpType.W)
  val SRLW = 13.asUInt(fuOpType.W)
  val SRAW = 14.asUInt(fuOpType.W)
  val COPY1= 15.asUInt(fuOpType.W)
}

class ALU extends Module {
  val io = IO(new Bundle(){
    val in = Flipped(Decoupled(new Bundle {
      val op1 = Output(UInt(XLEN.W))
      val op2 = Output(UInt(XLEN.W))
      val fun = Output(UInt(ALU_X.getWidth.W))
    }))
    val out = Decoupled(Output(UInt(XLEN.W)))
  })

  io.in.ready := true.B

  val (valid, src, src2, alu_fun) = (io.in.valid, io.in.bits.op1, io.in.bits.op2, io.in.bits.fun)
  def access(valid: Bool, src1: UInt, src2: UInt, func: UInt): UInt = {
    this.valid := valid
    this.src := src1
    this.src2 := src2
    this.alu_fun := func
    io.out.bits
  }

  val alu_out = Wire(UInt(XLEN.W))
  val alu_shamt = src2(5, 0)
  // alu operation
  alu_out := MuxCase(0.U, Array(
    (alu_fun === ALUOpType.ADD || alu_fun === ALUOpType.ADDW)   ->  (src + src2).asUInt(),
    (alu_fun === ALUOpType.SUB || alu_fun === ALUOpType.SUBW)   ->  (src - src2).asUInt(),
    (alu_fun === ALUOpType.AND)                           ->  (src & src2).asUInt(),
    (alu_fun === ALUOpType.OR)                            ->  (src | src2).asUInt(),
    (alu_fun === ALUOpType.XOR)                           ->  (src ^ src2).asUInt(),
    (alu_fun === ALUOpType.SLT)                           ->  (src.asSInt() < src2.asSInt()).asUInt(),
    (alu_fun === ALUOpType.SLTU)                          ->  (src < src2).asUInt(),
    (alu_fun === ALUOpType.SLL)                           ->  (src << alu_shamt).asUInt(),
    (alu_fun === ALUOpType.SLLW)                          ->  (src(31, 0) << alu_shamt(4, 0)).asUInt(),
    (alu_fun === ALUOpType.SRA)                           ->  (src.asSInt() >> alu_shamt).asUInt(),
    (alu_fun === ALUOpType.SRAW)                          ->  (src(31, 0).asSInt() >> alu_shamt(4, 0)).asUInt(),
    (alu_fun === ALUOpType.SRL)                           ->  (src >> alu_shamt).asUInt(),
    (alu_fun === ALUOpType.SRLW)                          ->  (src(31, 0).asUInt >> alu_shamt(4, 0)).asUInt(),
    (alu_fun === ALUOpType.COPY1)                         ->   src,

  ))

  io.out.valid := io.in.valid
  io.out.bits := MuxCase(alu_out, Array(
    (alu_fun === ALUOpType.ADDW)  ->   Cat(Fill(32, alu_out(31)), alu_out(31, 0)),
    (alu_fun === ALUOpType.SUBW)  ->   Cat(Fill(32, alu_out(31)), alu_out(31, 0)),
    (alu_fun === ALUOpType.SLLW)  ->   Cat(Fill(32, alu_out(31)), alu_out(31, 0)),
    (alu_fun === ALUOpType.SRAW)  ->   Cat(Fill(32, alu_out(31)), alu_out(31, 0)),
    (alu_fun === ALUOpType.SRLW)  ->   Cat(Fill(32, alu_out(31)), alu_out(31, 0)),
  ))
}
